#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

import numpy as np
from math import atan,atan2,pi

import robot
from robot import trotx,troty,trotz,rotx,roty,rotz,r2t,Link
from nao_defines import *


def RotXYZMatrix(xAngle,yAngle,zAngle):
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rx*Ry*Rz
    return r2t(R)

def RotZYXMatrix(zAngle,yAngle,xAngle):
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rz*Ry*Rx
    return r2t(R)

class lmRobot(robot.Robot):
    def __init__(self,name=""):
        super(lmRobot, self).__init__(name=name)

    def fkine(self,q):
        return robot.fkine(self,q)

    def ikine(self,tr):
        pass

    def jacob0(self,q):
        return robot.jacob0(self,q)

    def jacobn(self,q):
        return robot.jacobn(self,q)

    def solve0(self,tr,q,mask,ilimit=1000,stol=1e-6, alpha=1.0,norm_func=None,verstep=False):
        # mask 中的为 0 坐标系的移动和转动自由度, 即基坐标系
        from numpy.linalg import pinv,norm
        from numpy import array
        from robot import t2r
        import cv2
        if isinstance(q,list): q = np.array(q)
        if isinstance(tr,np.ndarray): tr = np.matrix(tr)
        count = 0 # 迭代次数
        inf = np.inf
        e = np.matrix([inf]*6) # 误差向量
        e = e.transpose()
        eprev = e.copy() # 前一个误差向量
        nm = inf # 误差值
        save_q = q.copy()
        save_e = e.copy()
        if norm_func is None: norm_func = norm

        while nm > stol:
            count += 1
            if count > ilimit:
                print("超出迭代限制")
                print("迭代次数:",count)
                print("误差:",nm)
                return q
            tq = self.fkine(q)
            e[0:3,0] = tr[0:3,3] - tq[0:3,3]
            rq = t2r(tq)
            e[3:6,0] = cv2.Rodrigues(rq.transpose()*t2r(tr))[0]
            if verstep:
                if norm_func(e[mask,0]) < norm_func(save_e[mask,0]):
                    # 更好的e则扩大步长,变步长 方法
                    save_q = q.copy()
                    save_e = e.copy()
                    alpha *= 2**(1/8.0)
                else:
                    # 使用保存的e,并缩短步长
                    q = save_q.copy()
                    e = save_e.copy()
                    alpha *= 0.5
            j = self.jacob0(q)
            #dq = j[m,:]*e[m,0]
            dq = pinv(j[mask,:])*e[mask,0]
            dq = alpha*dq
            q += array(dq).flatten()
            # wrap angles for revolute joints
            k = q > pi
            q[k] = q[k] - 2*pi
            k = q < -pi
            q[k] = q[k] + 2*pi
            nm = norm_func(e[mask,0])
            if norm_func(e[mask,0]) > 2*norm_func(eprev[mask,0]) :
                print("警告, ikine 发散")
            eprev = e
        print("求解成功,q:",q)
        print("迭代次数:",count)
        return q

    def solven(self,tr,q,mask,ilimit=1000,stol=1e-6, alpha=1.0,norm_func=None,verstep=False):
        # mask 中的为 n 坐标系的移动和转动自由度,即末端工具坐标系
        from numpy.linalg import inv,pinv,norm
        from numpy import array
        from robot import t2r
        import cv2
        if isinstance(q,list): q = np.array(q)
        if isinstance(tr,np.ndarray): tr = np.matrix(tr)
        count = 0 # 迭代次数
        inf = np.inf
        e = np.matrix([inf]*6) # 误差向量
        e = e.transpose()
        eprev = e.copy() # 前一个误差向量
        nm = inf # 误差值
        save_q = q.copy()
        save_e = e.copy()
        invtr = inv(tr) # 向末端执行器变换的矩阵
        tr = np.matrix(np.eye(4))

        if norm_func is None: norm_func = norm

        while nm > stol:
            count += 1
            if count > ilimit:
                print("超出迭代限制")
                print("迭代次数:",count)
                print("误差:",nm)
                return q
            tq = self.fkine(q)
            tq = invtr*tq # 在目标末端坐标系下的坐标矩阵
            e[0:3,0] = - tq[0:3,3]
            rq = t2r(tq)
            e[3:6,0] = cv2.Rodrigues(rq.transpose())[0]
            if verstep:
                if norm_func(e[mask,0]) < norm_func(save_e[mask,0]):
                    # 更好的e则扩大步长,变步长 方法
                    save_q = q.copy()
                    save_e = e.copy()
                    alpha *= 2**(1/8.0)
                else:
                    # 使用保存的e,并缩短步长
                    q = save_q.copy()
                    e = save_e.copy()
                    alpha *= 0.5
            j = self.jacobn(q)
            #dq = j[m,:]*e[m,0]
            dq = pinv(j[mask,:])*e[mask,0]
            dq = alpha*dq
            q += array(dq).flatten()
            # wrap angles for revolute joints
            k = q > pi
            q[k] = q[k] - 2*pi
            k = q < -pi
            q[k] = q[k] + 2*pi
            nm = norm_func(e[mask,0])
            if norm_func(e[mask,0]) > 2*norm_func(eprev[mask,0]) :
                print("警告, ikine 发散")
            eprev = e
        print("求解成功,q:",q)
        print("迭代次数:",count)
        return q


class LeftHand(lmRobot):
    def __init__(self,name="LeftHand"):
        super(LeftHand, self).__init__(name=name)
        self.base[1,3] = ShoulderOffsetY
        self.base[2,3] = ShoulderOffsetZ
        L = []
        L.append( Link(A=0,            alpha=-pi/2, D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# LShoulderPitch
        L.append( Link(A=0,            alpha=pi/2,  D=0, theta=pi/2, sigma=0, convention=Link.LINK_MDH) )# LShoulderRoll
        L.append( Link(A=ElbowOffsetY, alpha=pi/2,  D=UpperArmLength, theta=0, sigma=0, convention=Link.LINK_MDH) )# LElbowYaw
        L.append( Link(A=0,            alpha=-pi/2, D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# LElbowRoll
        L.append( Link(A=0,            alpha=pi/2,  D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# LWristYaw
        self.links=L
        self.tool[0,3] = LowerArmLength + HandOffsetX
        self.tool[2,3] = -HandOffsetZ
        self.tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.tool

    def lock(self,q,n=-1):
        '''
        锁住最后一个关节,使机器人退化
        :return:
        '''
        if n == -1:
            self.locked_n = len(self.links) - 1
        else:
            self.locked_n = n
        self.lockedL = self.links.pop(self.locked_n)
        self.before_locked_tool = self.tool.copy()
        self.tool = self.lockedL.tr(q)*self.before_locked_tool

    def unlock(self):
        self.links.insert(self.locked_n,self.lockedL)
        self.tool = self.before_locked_tool.copy()

    def ikine(self,tr):
        # 解析逆解
        inv = np.linalg.inv
        norm = np.linalg.norm
        base = self.base
        tool = self.tool
        L1 = self.links[0]
        L2 = self.links[1]
        L3 = self.links[2]
        T15 = inv(tool*inv(tr)*base)
        theta1 = atan2(-T15[2,3],T15[0,3])
        T25 = inv(L1.tr(theta1))*T15
        l1 = UpperArmLength
        l2 = ElbowOffsetY
        r14 = T25[0,3]
        r34 = T25[2,3]
        theta2 = atan2(l1*r34-l2*r14,l1*r14+l2*r34)
        T35 = inv(L2.tr(theta2))*T25
        theta3 = atan2(-T35[2,2],-T35[0,2])
        T45 = inv(L3.tr(theta3))*T35
        theta4 = atan2(T45[0,2],T45[2,2])
        theta5 = atan2(T45[1,0],T45[1,1])
        q = [theta1,theta2,theta3,theta4,theta5]
        err = norm(tr - self.fkine(q))
        return q,err

    def ikine2(self,tr,q0,mask=(0,1,2,3,5)): # 放弃绕末端执行器的y轴方向转动)
        def mynorm(e):
            e2 = e.copy()
            e2[0:3,0] *= 0.001
            return np.linalg.norm(e[:,0])
        return self.solven(tr,q0,mask=mask,verstep=True,norm_func=mynorm)

class RightHand(lmRobot):
    def __init__(self,name="RightHand"):
        super(RightHand, self).__init__(name=name)
        self.base[1,3] = -ShoulderOffsetY # 与LArm不同
        self.base[2,3] = ShoulderOffsetZ
        L = []
        L.append( Link(A=0,            alpha=-pi/2, D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# RShoulderPitch
        L.append( Link(A=0,            alpha=pi/2,  D=0, theta=pi/2, sigma=0, convention=Link.LINK_MDH) )# RShoulderRoll
        L.append( Link(A=-ElbowOffsetY, alpha=pi/2,  D=UpperArmLength, theta=0, sigma=0, convention=Link.LINK_MDH) )# RElbowYaw 与LArm不同
        L.append( Link(A=0,            alpha=-pi/2, D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# RElbowRoll
        L.append( Link(A=0,            alpha=pi/2,  D=0, theta=0,    sigma=0, convention=Link.LINK_MDH) )# RWristYaw
        self.links=L
        self.tool[0,3] = LowerArmLength + HandOffsetX
        self.tool[2,3] = -HandOffsetZ
        self.tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.tool

    def ikine(self,tr):
        inv = np.linalg.inv
        norm = np.linalg.norm
        base = self.base
        tool = self.tool
        L1 = self.links[0]
        L2 = self.links[1]
        L3 = self.links[2]
        T15 = inv(tool*inv(tr)*base)
        theta1 = atan2(-T15[2,3],T15[0,3])
        T25 = inv(L1.tr(theta1))*T15
        l1 = UpperArmLength
        l2 = ElbowOffsetY
        r14 = T25[0,3]
        r34 = T25[2,3]
        theta2 = atan2(l1*r34+l2*r14,l1*r14-l2*r34) # 与 Left 不同
        T35 = inv(L2.tr(theta2))*T25
        theta3 = atan2(T35[2,2],T35[0,2]) # 与 Left 不同
        T45 = inv(L3.tr(theta3))*T35
        theta4 = atan2(T45[0,2],T45[2,2])
        theta5 = atan2(T45[1,0],T45[1,1])
        q = [theta1,theta2,theta3,theta4,theta5]
        err = norm(tr - self.fkine(q))
        return q,err

    def ikine2(self,tr,q0,mask=(0,1,2,3,5)): # 放弃绕末端执行器的y轴方向转动)
        def mynorm(e):
            e2 = e.copy()
            e2[0:3,0] *= 0.001
            return np.linalg.norm(e[:,0])
        return self.solven(tr,q0,mask=mask,verstep=True,norm_func=mynorm)

class Head(lmRobot):
    def __init__(self,name):
        super(Head, self).__init__(name=name)
        self.base[2,3] = NeckOffsetZ
        L = []
        L.append( Link(A=0,alpha=0,      D=0,  theta=0,    sigma=0, convention=Link.LINK_MDH) ) # HeadYaw
        L.append( Link(A=0,alpha=-pi/2,  D=0, theta=-pi/2, sigma=0, convention=Link.LINK_MDH) ) # HeadPitch
        self.links=L
        self.tool = RotXYZMatrix(pi/2,pi/2,0)

class TopCamera(Head):
    def __init__(self,name):
        super(TopCamera, self).__init__(name=name)
        top_tool = np.matrix(np.eye(4))
        top_tool = troty(CameraTopAngle)*top_tool
        top_tool[0,3] = CameraTopX
        top_tool[2,3] = CameraTopZ
        top_tool = self.tool*top_tool
        self.tool = top_tool

class BottomCamera(Head):
    def __init__(self,name):
        super(BottomCamera, self).__init__(name=name)
        bottom_tool = np.matrix(np.eye(4))
        bottom_tool = troty(CameraBottomAngle)*bottom_tool
        bottom_tool[0,3] = CameraBottomX
        bottom_tool[2,3] = CameraBottomZ
        bottom_tool = self.tool*bottom_tool
        self.tool = bottom_tool

class Nao:
    def __init__(self,name):
        self.name = name
        self._loaded_camera_tool = False
        self.Head = Head(name=name+"Head")
        self.TopCamera = TopCamera(name=name+"_TopCamera")
        self.BottomCamera = BottomCamera(name=name+"_BottomCamera")
        self.LeftHand = LeftHand(name=name+"_LeftHand")
        self.RightHand = RightHand(name=name+"_RightHand")

    def load_camera_tool(self):
        import sql
        nao = sql.Nao.get_by_name(self.name)
        tool = nao.top_cam.tool
        if tool is not None:
            self.TopCamera.tool = tool
        tool = nao.bottom_cam.tool
        if tool is not None:
            self.BottomCamera.tool = tool
        self._loaded_camera_tool = True

nao = Nao(name="nao") # 通用求解正逆运动学
